from .UDP消息 import UDP消息, UDP消息类型
from .UDP协议 import UDP协议
from .UDP会话 import UDP会话
from .异常 import *
from .ID生成 import 循环ID生成器
from typing import Dict
import asyncio


class UDP服务器:
    def __init__(self, 协议: UDP协议, 客户端处理函数) -> None:
        self.协议 = 协议
        self.客户端处理函数 = 客户端处理函数
        self.会话ID生成器 = 循环ID生成器()
        self.会话: Dict[int, UDP会话] = {}
        self.loop = asyncio.get_running_loop()
        self.消息收取任务 = self.loop.create_task(self.收消息任务())

    async def 收消息任务(self):
        while not self.协议.链接关闭:
            data = await self.协议.read()
            if not data:
                raise UDP链接已关闭()
            消息 = UDP消息.拆包(*data)
            if 消息.消息类型 == UDP消息类型.请求建立链接:
                会话ID = self.会话ID生成器(self.会话)
            else:
                会话ID = 消息.会话ID
            会话 = self.会话.get(会话ID)
            if 会话 is None:
                会话 = UDP会话(self, self.协议, 会话ID, 消息.地址)
                self.会话[会话ID] = 会话
                会话.处理任务 = self.loop.create_task(self.客户端处理函数(会话))
            try:
                self.loop.create_task(会话.收消息(消息))
            except UDP链接已关闭:
                return
            except:
                self.关闭会话(会话ID)
                print('会话关闭')

    def 关闭会话(self, 会话ID: int):
        try:
            del self.会话[会话ID]
        except:
            pass

    async def serve_forever(self):
        while not self.协议.链接关闭:
            关闭的会话 = []
            for 会话 in self.会话.values():
                if 会话.会话关闭:
                    关闭的会话.append(会话.会话ID)
            if len(关闭的会话) > 0:
                print(关闭的会话)
            for 关闭的会话ID in 关闭的会话:
                self.关闭会话(关闭的会话ID)
            await asyncio.sleep(0.1)


async def start_server(客户端处理函数, 监听地址: str, 监听端口: int):
    loop = asyncio.get_running_loop()
    传输, 协议 = await loop.create_datagram_endpoint(
        lambda: UDP协议(),
        local_addr=(监听地址, 监听端口)
    )
    return UDP服务器(协议, 客户端处理函数)
